//
// Created by 周辉 on 2021/8/18.
//

#include <iostream>
#include <chrono>
#include <cmath>
#include "cuda_utils.h"
#include "../utils/logging.h"
#include "common.h"
#include "../utils/utils.h"
#include "yolov5.h"
#include "calibrator.h"
#include "../utils/dataset.h"

using namespace std;

namespace YoloV5
{
    static const int INPUT_H = Yolo::INPUT_H;
    static const int INPUT_W = Yolo::INPUT_W;
    static const int CLASS_NUM = Yolo::CLASS_NUM;
    static const int OUTPUT_SIZE = Yolo::MAX_OUTPUT_BBOX_COUNT * sizeof(Yolo::Detection) / sizeof(float) + 1;  // we assume the yololayer outputs no more than MAX_OUTPUT_BBOX_COUNT boxes that conf >= 0.1
    const char* INPUT_BLOB_NAME = "data";
    const char* OUTPUT_BLOB_NAME = "prob";
    static Logger gLogger;
    static int get_width(int x,float gw,int divisor=8){
	return int(ceil((x*gw)/divisor))*divisor;}
    
    static int get_depth(int x,float gd){
        if(x==1) return 1;
        int r=round(x*gd);
        if(x*gd-int(x*gd)==0.5&&(int(x*gd)%2)==0){
            --r;}
        return std::max<int>(r,1);
    }
    ICudaEngine* build_engine(unsigned int maxBatchSize, IBuilder* builder, IBuilderConfig* config, DataType dt, float& gd, float& gw, std::string& wts_name) {
        INetworkDefinition* network = builder->createNetworkV2(0U);

        // Create input tensor of shape {3, INPUT_H, INPUT_W} with name INPUT_BLOB_NAME
        ITensor* data = network->addInput(INPUT_BLOB_NAME, dt, Dims3{ 3, INPUT_H, INPUT_W });
        assert(data);

        std::map<std::string, Weights> weightMap = loadWeights(wts_name);

        /* ------ yolov5 backbone------ */
        auto focus0 = focus(network, weightMap, *data, 3, get_width(64, gw), 3, "model.0");
        auto conv1 = convBlock(network, weightMap, *focus0->getOutput(0), get_width(128, gw), 3, 2, 1, "model.1");
        auto bottleneck_CSP2 = C3(network, weightMap, *conv1->getOutput(0), get_width(128, gw), get_width(128, gw), get_depth(3, gd), true, 1, 0.5, "model.2");
        auto conv3 = convBlock(network, weightMap, *bottleneck_CSP2->getOutput(0), get_width(256, gw), 3, 2, 1, "model.3");
        auto bottleneck_csp4 = C3(network, weightMap, *conv3->getOutput(0), get_width(256, gw), get_width(256, gw), get_depth(9, gd), true, 1, 0.5, "model.4");
        auto conv5 = convBlock(network, weightMap, *bottleneck_csp4->getOutput(0), get_width(512, gw), 3, 2, 1, "model.5");
        auto bottleneck_csp6 = C3(network, weightMap, *conv5->getOutput(0), get_width(512, gw), get_width(512, gw), get_depth(9, gd), true, 1, 0.5, "model.6");
        auto conv7 = convBlock(network, weightMap, *bottleneck_csp6->getOutput(0), get_width(1024, gw), 3, 2, 1, "model.7");
        auto spp8 = SPP(network, weightMap, *conv7->getOutput(0), get_width(1024, gw), get_width(1024, gw), 5, 9, 13, "model.8");

        /* ------ yolov5 head ------ */
        auto bottleneck_csp9 = C3(network, weightMap, *spp8->getOutput(0), get_width(1024, gw), get_width(1024, gw), get_depth(3, gd), false, 1, 0.5, "model.9");
        auto conv10 = convBlock(network, weightMap, *bottleneck_csp9->getOutput(0), get_width(512, gw), 1, 1, 1, "model.10");

        auto upsample11 = network->addResize(*conv10->getOutput(0));
        assert(upsample11);
        upsample11->setResizeMode(ResizeMode::kNEAREST);
        upsample11->setOutputDimensions(bottleneck_csp6->getOutput(0)->getDimensions());

        ITensor* inputTensors12[] = { upsample11->getOutput(0), bottleneck_csp6->getOutput(0) };
        auto cat12 = network->addConcatenation(inputTensors12, 2);
        auto bottleneck_csp13 = C3(network, weightMap, *cat12->getOutput(0), get_width(1024, gw), get_width(512, gw), get_depth(3, gd), false, 1, 0.5, "model.13");
        auto conv14 = convBlock(network, weightMap, *bottleneck_csp13->getOutput(0), get_width(256, gw), 1, 1, 1, "model.14");

        auto upsample15 = network->addResize(*conv14->getOutput(0));
        assert(upsample15);
        upsample15->setResizeMode(ResizeMode::kNEAREST);
        upsample15->setOutputDimensions(bottleneck_csp4->getOutput(0)->getDimensions());

        ITensor* inputTensors16[] = { upsample15->getOutput(0), bottleneck_csp4->getOutput(0) };
        auto cat16 = network->addConcatenation(inputTensors16, 2);

        auto bottleneck_csp17 = C3(network, weightMap, *cat16->getOutput(0), get_width(512, gw), get_width(256, gw), get_depth(3, gd), false, 1, 0.5, "model.17");

        /* ------ detect ------ */
        IConvolutionLayer* det0 = network->addConvolutionNd(*bottleneck_csp17->getOutput(0), 3 * (Yolo::CLASS_NUM + 5), DimsHW{ 1, 1 }, weightMap["model.24.m.0.weight"], weightMap["model.24.m.0.bias"]);
        auto conv18 = convBlock(network, weightMap, *bottleneck_csp17->getOutput(0), get_width(256, gw), 3, 2, 1, "model.18");
        ITensor* inputTensors19[] = { conv18->getOutput(0), conv14->getOutput(0) };
        auto cat19 = network->addConcatenation(inputTensors19, 2);
        auto bottleneck_csp20 = C3(network, weightMap, *cat19->getOutput(0), get_width(512, gw), get_width(512, gw), get_depth(3, gd), false, 1, 0.5, "model.20");
        IConvolutionLayer* det1 = network->addConvolutionNd(*bottleneck_csp20->getOutput(0), 3 * (Yolo::CLASS_NUM + 5), DimsHW{ 1, 1 }, weightMap["model.24.m.1.weight"], weightMap["model.24.m.1.bias"]);
        auto conv21 = convBlock(network, weightMap, *bottleneck_csp20->getOutput(0), get_width(512, gw), 3, 2, 1, "model.21");
        ITensor* inputTensors22[] = { conv21->getOutput(0), conv10->getOutput(0) };
        auto cat22 = network->addConcatenation(inputTensors22, 2);
        auto bottleneck_csp23 = C3(network, weightMap, *cat22->getOutput(0), get_width(1024, gw), get_width(1024, gw), get_depth(3, gd), false, 1, 0.5, "model.23");
        IConvolutionLayer* det2 = network->addConvolutionNd(*bottleneck_csp23->getOutput(0), 3 * (Yolo::CLASS_NUM + 5), DimsHW{ 1, 1 }, weightMap["model.24.m.2.weight"], weightMap["model.24.m.2.bias"]);

        auto yolo = addYoLoLayer(network, weightMap, "model.24", std::vector<IConvolutionLayer*>{det0, det1, det2});
        yolo->getOutput(0)->setName(OUTPUT_BLOB_NAME);
        network->markOutput(*yolo->getOutput(0));

        // Build engine
        builder->setMaxBatchSize(maxBatchSize);
        config->setMaxWorkspaceSize(16 * (1 << 20));  // 16MB
    #if defined(USE_FP16)
            config->setFlag(BuilderFlag::kFP16);
    #elif defined(USE_INT8)
            std::cout << "Your platform support int8: " << (builder->platformHasFastInt8() ? "true" : "false") << std::endl;
        assert(builder->platformHasFastInt8());
        config->setFlag(BuilderFlag::kINT8);
        Int8EntropyCalibrator2* calibrator = new Int8EntropyCalibrator2(1, INPUT_W, INPUT_H, "./coco_calib/", "int8calib.table", INPUT_BLOB_NAME);
        config->setInt8Calibrator(calibrator);
    #endif

        std::cout << "Building engine, please wait for a while..." << std::endl;
        ICudaEngine* engine = builder->buildEngineWithConfig(*network, *config);
        std::cout << "Build engine successfully!" << std::endl;

        // Don't need the network any more
        network->destroy();

        // Release host memory
        for (auto& mem : weightMap)
        {
            free((void*)(mem.second.values));
        }

        return engine;
    }

    ICudaEngine* build_engine_p6(unsigned int maxBatchSize, IBuilder* builder, IBuilderConfig* config, DataType dt, float& gd, float& gw, std::string& wts_name) {
        INetworkDefinition* network = builder->createNetworkV2(0U);

        // Create input tensor of shape {3, INPUT_H, INPUT_W} with name INPUT_BLOB_NAME
        ITensor* data = network->addInput(INPUT_BLOB_NAME, dt, Dims3{ 3, INPUT_H, INPUT_W });
        assert(data);

        std::map<std::string, Weights> weightMap = loadWeights(wts_name);

        /* ------ yolov5 backbone------ */
        auto focus0 = focus(network, weightMap, *data, 3, get_width(64, gw), 3, "model.0");
        auto conv1 = convBlock(network, weightMap, *focus0->getOutput(0), get_width(128, gw), 3, 2, 1, "model.1");
        auto c3_2 = C3(network, weightMap, *conv1->getOutput(0), get_width(128, gw), get_width(128, gw), get_depth(3, gd), true, 1, 0.5, "model.2");
        auto conv3 = convBlock(network, weightMap, *c3_2->getOutput(0), get_width(256, gw), 3, 2, 1, "model.3");
        auto c3_4 = C3(network, weightMap, *conv3->getOutput(0), get_width(256, gw), get_width(256, gw), get_depth(9, gd), true, 1, 0.5, "model.4");
        auto conv5 = convBlock(network, weightMap, *c3_4->getOutput(0), get_width(512, gw), 3, 2, 1, "model.5");
        auto c3_6 = C3(network, weightMap, *conv5->getOutput(0), get_width(512, gw), get_width(512, gw), get_depth(9, gd), true, 1, 0.5, "model.6");
        auto conv7 = convBlock(network, weightMap, *c3_6->getOutput(0), get_width(768, gw), 3, 2, 1, "model.7");
        auto c3_8 = C3(network, weightMap, *conv7->getOutput(0), get_width(768, gw), get_width(768, gw), get_depth(3, gd), true, 1, 0.5, "model.8");
        auto conv9 = convBlock(network, weightMap, *c3_8->getOutput(0), get_width(1024, gw), 3, 2, 1, "model.9");
        auto spp10 = SPP(network, weightMap, *conv9->getOutput(0), get_width(1024, gw), get_width(1024, gw), 3, 5, 7, "model.10");
        auto c3_11 = C3(network, weightMap, *spp10->getOutput(0), get_width(1024, gw), get_width(1024, gw), get_depth(3, gd), false, 1, 0.5, "model.11");

        /* ------ yolov5 head ------ */
        auto conv12 = convBlock(network, weightMap, *c3_11->getOutput(0), get_width(768, gw), 1, 1, 1, "model.12");
        auto upsample13 = network->addResize(*conv12->getOutput(0));
        assert(upsample13);
        upsample13->setResizeMode(ResizeMode::kNEAREST);
        upsample13->setOutputDimensions(c3_8->getOutput(0)->getDimensions());
        ITensor* inputTensors14[] = { upsample13->getOutput(0), c3_8->getOutput(0) };
        auto cat14 = network->addConcatenation(inputTensors14, 2);
        auto c3_15 = C3(network, weightMap, *cat14->getOutput(0), get_width(1536, gw), get_width(768, gw), get_depth(3, gd), false, 1, 0.5, "model.15");

        auto conv16 = convBlock(network, weightMap, *c3_15->getOutput(0), get_width(512, gw), 1, 1, 1, "model.16");
        auto upsample17 = network->addResize(*conv16->getOutput(0));
        assert(upsample17);
        upsample17->setResizeMode(ResizeMode::kNEAREST);
        upsample17->setOutputDimensions(c3_6->getOutput(0)->getDimensions());
        ITensor* inputTensors18[] = { upsample17->getOutput(0), c3_6->getOutput(0) };
        auto cat18 = network->addConcatenation(inputTensors18, 2);
        auto c3_19 = C3(network, weightMap, *cat18->getOutput(0), get_width(1024, gw), get_width(512, gw), get_depth(3, gd), false, 1, 0.5, "model.19");

        auto conv20 = convBlock(network, weightMap, *c3_19->getOutput(0), get_width(256, gw), 1, 1, 1, "model.20");
        auto upsample21 = network->addResize(*conv20->getOutput(0));
        assert(upsample21);
        upsample21->setResizeMode(ResizeMode::kNEAREST);
        upsample21->setOutputDimensions(c3_4->getOutput(0)->getDimensions());
        ITensor* inputTensors21[] = { upsample21->getOutput(0), c3_4->getOutput(0) };
        auto cat22 = network->addConcatenation(inputTensors21, 2);
        auto c3_23 = C3(network, weightMap, *cat22->getOutput(0), get_width(512, gw), get_width(256, gw), get_depth(3, gd), false, 1, 0.5, "model.23");

        auto conv24 = convBlock(network, weightMap, *c3_23->getOutput(0), get_width(256, gw), 3, 2, 1, "model.24");
        ITensor* inputTensors25[] = { conv24->getOutput(0), conv20->getOutput(0) };
        auto cat25 = network->addConcatenation(inputTensors25, 2);
        auto c3_26 = C3(network, weightMap, *cat25->getOutput(0), get_width(1024, gw), get_width(512, gw), get_depth(3, gd), false, 1, 0.5, "model.26");

        auto conv27 = convBlock(network, weightMap, *c3_26->getOutput(0), get_width(512, gw), 3, 2, 1, "model.27");
        ITensor* inputTensors28[] = { conv27->getOutput(0), conv16->getOutput(0) };
        auto cat28 = network->addConcatenation(inputTensors28, 2);
        auto c3_29 = C3(network, weightMap, *cat28->getOutput(0), get_width(1536, gw), get_width(768, gw), get_depth(3, gd), false, 1, 0.5, "model.29");

        auto conv30 = convBlock(network, weightMap, *c3_29->getOutput(0), get_width(768, gw), 3, 2, 1, "model.30");
        ITensor* inputTensors31[] = { conv30->getOutput(0), conv12->getOutput(0) };
        auto cat31 = network->addConcatenation(inputTensors31, 2);
        auto c3_32 = C3(network, weightMap, *cat31->getOutput(0), get_width(2048, gw), get_width(1024, gw), get_depth(3, gd), false, 1, 0.5, "model.32");

        /* ------ detect ------ */
        IConvolutionLayer* det0 = network->addConvolutionNd(*c3_23->getOutput(0), 3 * (Yolo::CLASS_NUM + 5), DimsHW{ 1, 1 }, weightMap["model.33.m.0.weight"], weightMap["model.33.m.0.bias"]);
        IConvolutionLayer* det1 = network->addConvolutionNd(*c3_26->getOutput(0), 3 * (Yolo::CLASS_NUM + 5), DimsHW{ 1, 1 }, weightMap["model.33.m.1.weight"], weightMap["model.33.m.1.bias"]);
        IConvolutionLayer* det2 = network->addConvolutionNd(*c3_29->getOutput(0), 3 * (Yolo::CLASS_NUM + 5), DimsHW{ 1, 1 }, weightMap["model.33.m.2.weight"], weightMap["model.33.m.2.bias"]);
        IConvolutionLayer* det3 = network->addConvolutionNd(*c3_32->getOutput(0), 3 * (Yolo::CLASS_NUM + 5), DimsHW{ 1, 1 }, weightMap["model.33.m.3.weight"], weightMap["model.33.m.3.bias"]);

        auto yolo = addYoLoLayer(network, weightMap, "model.33", std::vector<IConvolutionLayer*>{det0, det1, det2, det3});
        yolo->getOutput(0)->setName(OUTPUT_BLOB_NAME);
        network->markOutput(*yolo->getOutput(0));

        // Build engine
        builder->setMaxBatchSize(maxBatchSize);
        config->setMaxWorkspaceSize(16 * (1 << 20));  // 16MB
    #if defined(USE_FP16)
            config->setFlag(BuilderFlag::kFP16);
    #elif defined(USE_INT8)
            std::cout << "Your platform support int8: " << (builder->platformHasFastInt8() ? "true" : "false") << std::endl;
        assert(builder->platformHasFastInt8());
        config->setFlag(BuilderFlag::kINT8);
        Int8EntropyCalibrator2* calibrator = new Int8EntropyCalibrator2(1, INPUT_W, INPUT_H, "./coco_calib/", "int8calib.table", INPUT_BLOB_NAME);
        config->setInt8Calibrator(calibrator);
    #endif

        std::cout << "Building engine, please wait for a while..." << std::endl;
        ICudaEngine* engine = builder->buildEngineWithConfig(*network, *config);
        std::cout << "Build engine successfully!" << std::endl;

        // Don't need the network any more
        network->destroy();

        // Release host memory
        for (auto& mem : weightMap)
        {
            free((void*)(mem.second.values));
        }

        return engine;
    }
    void APIToModel(unsigned int maxBatchSize, IHostMemory** modelStream, bool& is_p6, float& gd, float& gw, std::string& wts_name) {
        // Create builder
        IBuilder* builder = createInferBuilder(gLogger);
        IBuilderConfig* config = builder->createBuilderConfig();

        // Create model to populate the network, then set the outputs and create an engine
        ICudaEngine *engine = nullptr;
        if (is_p6) {
            engine = build_engine_p6(maxBatchSize, builder, config, DataType::kFLOAT, gd, gw, wts_name);
        } else {
            engine = build_engine(maxBatchSize, builder, config, DataType::kFLOAT, gd, gw, wts_name);
        }
        assert(engine != nullptr);

        // Serialize the engine
        (*modelStream) = engine->serialize();

        // Close everything down
        engine->destroy();
        builder->destroy();
        config->destroy();
    }

    void doInference(IExecutionContext &context, float *input, float *output, int batchSize) {
        const ICudaEngine &engine = context.getEngine();

        // Pointers to input and output device buffers to pass to engine.
        // Engine requires exactly IEngine::getNbBindings() number of buffers.
        assert(engine.getNbBindings() == 2);
        void *buffers[2];

        // In order to bind the buffers, we need to know the names of the input and output tensors.
        // Note that indices are guaranteed to be less than IEngine::getNbBindings()
        const int inputIndex = engine.getBindingIndex(INPUT_BLOB_NAME);
        const int outputIndex = engine.getBindingIndex(OUTPUT_BLOB_NAME);

        // Create GPU buffers on device
        CHECK(cudaMalloc(&buffers[inputIndex], batchSize * 3 * INPUT_H * INPUT_W * sizeof(float)));
        CHECK(cudaMalloc(&buffers[outputIndex], batchSize * OUTPUT_SIZE * sizeof(float)));

        // Create stream
        cudaStream_t stream;
        CHECK(cudaStreamCreate(&stream));

        // DMA input batch data to device, infer on the batch asynchronously, and DMA output back to host
        CHECK(cudaMemcpyAsync(buffers[inputIndex], input, batchSize * 3 * INPUT_H * INPUT_W * sizeof(float),
                              cudaMemcpyHostToDevice, stream));
        context.enqueue(batchSize, buffers, stream, nullptr);
        CHECK(cudaMemcpyAsync(output, buffers[outputIndex], batchSize * OUTPUT_SIZE * sizeof(float), cudaMemcpyDeviceToHost,
                              stream));
        cudaStreamSynchronize(stream);

        // Release stream and buffers
        cudaStreamDestroy(stream);
        CHECK(cudaFree(buffers[inputIndex]));
        CHECK(cudaFree(buffers[outputIndex]));
    }
    bool parse_args(int argc, char** argv, std::string& wts, std::string& engine, bool& is_p6, float& gd, float& gw, std::string& img_dir) {
        if (argc < 4) return false;
        if (std::string(argv[1]) == "-s" && (argc == 5 || argc == 7)) {
            wts = std::string(argv[2]);
            engine = std::string(argv[3]);
            auto net = std::string(argv[4]);
            if (net[0] == 's') {
                gd = 0.33;
                gw = 0.50;
            } else if (net[0] == 'm') {
                gd = 0.67;
                gw = 0.75;
            } else if (net[0] == 'l') {
                gd = 1.0;
                gw = 1.0;
            } else if (net[0] == 'x') {
                gd = 1.33;
                gw = 1.25;
            } else if (net[0] == 'c' && argc == 7) {
                gd = atof(argv[5]);
                gw = atof(argv[6]);
            } else {
                return false;
            }
            if (net.size() == 2 && net[1] == '6') {
                is_p6 = true;
            }
        } else if (std::string(argv[1]) == "-d" && argc == 4) {
            engine = std::string(argv[2]);
            img_dir = std::string(argv[3]);
        } else {
            return false;
        }
        return true;
    }
    IRuntime *runtime;
    ICudaEngine *engine;
    IExecutionContext *context;
    void InitYoloV5Engine()
    {
        char *trtModelStream{nullptr};
        size_t size{0};
        std::string engine_name = STR2(NET);

        // 测试使用
        if (BATCH_SIZE == 2)
        {
            engine_name = ENGIN_PATH + "yolov5" + engine_name + ".engine";
            cout << "使用 BATCH_SIZE = 2 的引擎: " << engine_name << endl;
        }
        else
        {
            engine_name = ENGIN_PATH + "yolov5" + engine_name + ".engine";
            cout << "ceshi使用 BATCH_SIZE = 1 的引擎: " << engine_name << endl;
        }

        std::ifstream file(engine_name, std::ios::binary);
        if (file.good())
        {
            file.seekg(0, file.end);
            size = file.tellg();
            file.seekg(0, file.beg);
            trtModelStream = new char[size];
            assert(trtModelStream);
            file.read(trtModelStream, size);
            file.close();
        }

        runtime = createInferRuntime(gLogger);
        assert(runtime != nullptr);
        engine = runtime->deserializeCudaEngine(trtModelStream, size);
        assert(engine != nullptr);
        context = engine->createExecutionContext();
        assert(context != nullptr);
        delete[] trtModelStream;
    }

    void ReleaseYoloV5Engine()
    {
        context->destroy();
        engine->destroy();
        runtime->destroy();
    }

    static float data[BATCH_SIZE * 3 * INPUT_H * INPUT_W];
    static float prob[BATCH_SIZE * OUTPUT_SIZE];

    vector<vector<Yolo::Detection>> AnalyzeBatch(vector<cv::Mat> &frames)
    {
        int batch_size = frames.size();

        for (int batch_index = 0; batch_index < batch_size; batch_index++)
        {
            if (frames[batch_index].empty())
            {
                continue;
            }
            cv::Mat pr_img = preprocess_img(frames[batch_index]);

            int pixel_pos = 0;
            for (int row = 0; row < INPUT_H; ++row)
            {
                uchar *uc_pixel = pr_img.data + row * pr_img.step;
                for (int col = 0; col < INPUT_W; ++col)
                {
                    data[batch_index * 3 * INPUT_H * INPUT_W + pixel_pos] = (float) uc_pixel[2] / 255.0;
                    data[batch_index * 3 * INPUT_H * INPUT_W + pixel_pos + INPUT_H * INPUT_W] = (float) uc_pixel[1] / 255.0;
                    data[batch_index * 3 * INPUT_H * INPUT_W + pixel_pos + 2 * INPUT_H * INPUT_W] =
                            (float) uc_pixel[0] / 255.0;
                    uc_pixel += 3;
                    ++pixel_pos;
                }
            }
        }
        // Run inference
        auto start = std::chrono::system_clock::now();
        doInference(*context, data, prob, batch_size);
        auto end = std::chrono::system_clock::now();
        std::cout << std::chrono::duration_cast<std::chrono::milliseconds>(end - start).count() << "ms" << std::endl;

        vector<vector<Yolo::Detection>> batch_results(batch_size);
        for (int batch_index = 0; batch_index < batch_size; batch_index++)
        {
            auto &res = batch_results[batch_index];
            nms(res, &prob[batch_index * OUTPUT_SIZE], CONF_THRESH, NMS_THRESH);

            cout << "batch[" << batch_index << "] 有目标数：" << res.size() << endl;
        }

        return batch_results;
    }

    vector<Yolo::Detection> AnalyzeOneShot(cv::Mat &frame)
    {
        int fcount = 0;
        int frame_count = 0;

        fcount++;
        frame_count++;

        for (int b = 0; b < fcount; b++)
        {
            if (frame.empty())
            {
                continue;
            }
            cv::Mat pr_img = preprocess_img(frame); // letterbox BGR to RGB

            int i = 0;
            for (int row = 0; row < INPUT_H; ++row)
            {
                uchar *uc_pixel = pr_img.data + row * pr_img.step;
                for (int col = 0; col < INPUT_W; ++col)
                {
                    data[b * 3 * INPUT_H * INPUT_W + i] = (float) uc_pixel[2] / 255.0;
                    data[b * 3 * INPUT_H * INPUT_W + i + INPUT_H * INPUT_W] = (float) uc_pixel[1] / 255.0;
                    data[b * 3 * INPUT_H * INPUT_W + i + 2 * INPUT_H * INPUT_W] = (float) uc_pixel[0] / 255.0;
                    uc_pixel += 3;
                    ++i;
                }
            }
        }

        // Run inference
        auto start = std::chrono::system_clock::now();
        doInference(*context, data, prob, 1);
        auto end = std::chrono::system_clock::now();
        std::cout << std::chrono::duration_cast<std::chrono::milliseconds>(end - start).count() << "ms" << std::endl;

        std::vector<Yolo::Detection> myres;
        nms(myres, &prob[0], CONF_THRESH, NMS_THRESH);

        std::vector<std::vector<Yolo::Detection>> batch_res(fcount);
        for (int b = 0; b < fcount; b++)
        {
            auto &res = batch_res[b];
            nms(res, &prob[b * OUTPUT_SIZE], CONF_THRESH, NMS_THRESH);
        }

        return batch_res[0];
    }

    void DrawYoloOutput(cv::Mat &frame, std::vector<Yolo::Detection> result)
    {
        for (size_t j = 0; j < result.size(); j++)
        {
            cv::Rect r = get_rect(frame, result[j].bbox);
            cv::rectangle(frame, r, cv::Scalar(0x27, 0xC1, 0x36), 2);
            cv::putText(frame, GetNameFromID(result[j].class_id), cv::Point(r.x, r.y - 1), cv::FONT_HERSHEY_PLAIN, 1.2,
                        cv::Scalar(0xFF, 0xFF, 0xFF), 2);
        }
    }

}
